import psycopg2
import psycopg2.extras
from datetime import datetime
import hashlib
import os


def write_report(schema_diff, data_diff, source_config, target_config, output_dir="diff_reports"):
    """生成结构化对比报告"""
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filepath = os.path.join(output_dir, f"diff_report_{timestamp}.txt")

    report = [
        "=" * 60,
        f"数据库对比报告 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "=" * 60,
        f"源数据库: {source_config['host']}:{source_config['port']}:{source_config['dbname']}",
        f"目标数据库: {target_config['host']}:{target_config['port']}:{target_config['dbname']}",
        "-" * 60
    ]

    report.append("\n[结构差异]")
    report.extend(schema_diff if schema_diff else ["无结构差异"])

    report.append("\n[数据差异]")
    report.extend(data_diff if data_diff else ["无数据差异"])

    try:
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write('\n'.join(report))
        return filepath
    except Exception as e:
        print(f"报告生成失败: {str(e)}")
        return None


class DatabaseComparator:
    def __init__(self, source_config, target_config, batch_size=1000, output_dir="diff_reports"):
        self.source_config = source_config
        self.target_config = target_config
        self.batch_size = batch_size
        self.output_dir = output_dir
        self.schema_diff = []
        self.data_diff = []
        self.source_conn = None
        self.target_conn = None

    def _connect(self):
        """建立数据库连接"""
        try:
            self.source_conn = psycopg2.connect(**self.source_config)
            self.target_conn = psycopg2.connect(**self.target_config)
            self.source_conn.autocommit = True
            self.target_conn.autocommit = True
        except Exception as e:
            raise ConnectionError(f"数据库连接失败: {str(e)}")

    def _get_table_metadata(self, conn):
        """获取完整的表结构元数据"""
        with conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cursor:
            cursor.execute("""
                SELECT table_name, column_name, data_type, 
                       is_nullable, column_default, ordinal_position
                FROM information_schema.columns
                WHERE table_schema = 'public'
                ORDER BY table_name, ordinal_position;
            """)
            return cursor.fetchall()

    def _compare_schemas(self):
        """执行表结构对比"""
        source_tables = self._get_table_metadata(self.source_conn)
        target_tables = self._get_table_metadata(self.target_conn)

        # 转换数据结构
        source_dict = {}
        for row in source_tables:
            source_dict.setdefault(row['table_name'], []).append(dict(row))

        target_dict = {}
        for row in target_tables:
            target_dict.setdefault(row['table_name'], []).append(dict(row))

        # 检查缺失表
        for table in set(source_dict) - set(target_dict):
            self.schema_diff.append(f"表缺失: {table} 在目标数据库不存在")
        for table in set(target_dict) - set(source_dict):
            self.schema_diff.append(f"表冗余: {table} 在源数据库不存在")

        # 对比共有表结构
        common_tables = set(source_dict) & set(target_dict)
        for table in common_tables:
            source_cols = {col['column_name']: col for col in source_dict[table]}
            target_cols = {col['column_name']: col for col in target_dict[table]}

            # 列缺失检查
            for col in source_cols:
                if col not in target_cols:
                    self.schema_diff.append(f"表 {table}: 列缺失 - {col}")
                    continue

                # 列属性对比
                for attr in ['data_type', 'is_nullable', 'column_default']:
                    if source_cols[col][attr] != target_cols[col][attr]:
                        self.schema_diff.append(
                            f"表 {table}.{col}: {attr} 不一致 | "
                            f"源: {source_cols[col][attr]} | 目标: {target_cols[col][attr]}"
                        )

            # 冗余列检查
            for col in set(target_cols) - set(source_cols):
                self.schema_diff.append(f"表 {table}: 列冗余 - {col}")

    def _get_primary_key(self, table):
        """获取表的主键列"""
        with self.source_conn.cursor() as cursor:
            cursor.execute(f"""
                SELECT a.attname
                FROM pg_index i
                JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                WHERE i.indrelid = %s::regclass AND i.indisprimary
            """, (table,))
            print("sql:", f"""
                SELECT a.attname
                FROM pg_index i
                JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
                WHERE i.indrelid = %s::regclass AND i.indisprimary
            """, (table,))
            result = cursor.fetchone()
            return result if result else 'ctid'

    def _compare_table_data(self, table):
        """分页对比表数据"""
        pkey = self._get_primary_key(table)
        print("table", table, "pkey", pkey[0])
        if pkey[0] == "c":
            print("table" + table + "没有主键，不做对比")
            self.data_diff.append(
                f"表 {table} 没有主键，不做对比\n"
            )
            return
        offset = 0
        mismatch_count = 0

        while True:
            # 分页查询
            with self.source_conn.cursor() as src_cur, self.target_conn.cursor() as tgt_cur:
                query = f"SELECT * FROM {table} ORDER BY {pkey[0]} LIMIT %s OFFSET %s"

                src_cur.execute(query, (self.batch_size, offset))
                source_batch = src_cur.fetchall()

                tgt_cur.execute(query, (self.batch_size, offset))
                target_batch = tgt_cur.fetchall()

                # 终止条件
                if not source_batch and not target_batch:
                    break

                # 快速哈希比对
                source_hash = hashlib.md5(str(source_batch).encode()).hexdigest()
                target_hash = hashlib.md5(str(target_batch).encode()).hexdigest()

                if source_hash != target_hash:
                    # 详细记录差异
                    for idx, (src_row, tgt_row) in enumerate(zip(source_batch, target_batch)):
                        if src_row != tgt_row:
                            mismatch_count += 1
                            self.data_diff.append(
                                f"表 {table} 第{offset + idx + 1}行不一致\n"
                                f"源: {src_row}\n目: {tgt_row}\n"
                            )

                offset += self.batch_size

        if mismatch_count > 0:
            self.data_diff.append(f"表 {table} 共发现 {mismatch_count} 处数据不一致")

    def compare(self):
        """执行完整对比流程"""
        try:
            self._connect()
            start_time = datetime.now()

            print("开始结构对比...")
            self._compare_schemas()

            print("开始数据对比...")
            source_tables = [t['table_name'] for t in self._get_table_metadata(self.source_conn)]
            target_tables = [t['table_name'] for t in self._get_table_metadata(self.target_conn)]
            common_tables = set(source_tables) & set(target_tables)

            for table in common_tables:
                print(f"正在对比表: {table}")
                if table in {"opt_nedication_detail_y2018m08"}:
                    print("该表不在对比范围" + table )
                    continue
                self._compare_table_data(table)

            # 生成报告
            report_path = write_report(
                self.schema_diff,
                self.data_diff,
                self.source_config,
                self.target_config,
                self.output_dir
            )

            # 输出统计信息
            total_time = datetime.now() - start_time
            print(f"\n对比完成，总耗时: {total_time}")
            print(f"发现结构差异: {len(self.schema_diff)} 处")
            print(f"发现数据差异: {len(self.data_diff)} 处")
            print(f"报告文件路径: {os.path.abspath(report_path)}")

        finally:
            if self.source_conn:
                self.source_conn.close()
            if self.target_conn:
                self.target_conn.close()


if __name__ == "__main__":
    # 数据库配置
    SOURCE_DB = {
        'host': '172.16.1.21',
        'dbname': 'upms',
        'user': 'postgres',
        'password': 'postgres',
        'port': 31013
    }

    TARGET_DB = {
        'host': '172.16.1.21',
        'dbname': 'upms',
        'user': 'postgres',
        'password': 'postgres',
        'port': 31003
    }

    # 初始化比较器
    comparator = DatabaseComparator(
        source_config=SOURCE_DB,
        target_config=TARGET_DB,
        batch_size=2000,
        output_dir="diff_reports"
    )

    # 执行对比
    comparator.compare()
